import numpy as np
import matplotlib.pyplot as plt
import time
import pandas as pd
import math as math

from scipy.interpolate import UnivariateSpline
from scipy.stats import gaussian_kde, beta, lognorm, norm, normaltest, kstest, ks_2samp,mannwhitneyu,ttest_ind




def create_localgridcoords(elev):
    irow=np.arange(elev.shape[1])
    icol=np.arange(elev.shape[0])
    return np.meshgrid(irow,icol[::-1])
    
def create_localxy(n,pixsize):
    icol=np.asarray([-(n/2)+i for i in range(n)])*pixsize
    irow=np.asarray([(n/2)-i for i in range(n)])*pixsize
    xy=np.meshgrid(icol,irow)
    return xy[0],xy[1]

def extract_localgrid(elev,cx,cy,n):
    irow=np.asarray([-(n/2)+i+cx for i in range(n)])
    icol=np.asarray([-(n/2)+i+cy for i in range(n)])
    locgrid=elev[irow[0]:irow[-1]+1,icol[0]:icol[-1]+1]
    if np.isnan(np.sum(locgrid))==False:
        return locgrid-locgrid[n/2,n/2]
    else:
        return np.array([])

def compute_gradient_aspect(dx,dy):
    grad=np.degrees(np.arctan(np.sqrt(dx**2+dy**2)))
    if grad<1:asp=np.nan
    else:
        if dx<0:asp=90-np.degrees(np.arctan(dy/dx))
        elif dx>0:asp=270-np.degrees(np.arctan(dy/dx))
        else:
            if dy<0:asp=0
            else:asp=180
    return np.round(grad,1), np.round(asp,1)
       


def generate_gradient_aspect(elev,n,pixsize):

    #creating matrix x(y)
    x,y=create_localxy(n,pixsize)
    sx2=np.sum(x**2)

    #initializong array for grad
    grad=np.nan*np.zeros(elev.shape)
    asp=np.nan*np.zeros(elev.shape)

    #loop through elev grid
    for i in range((n/2),elev.shape[0]-(n/2)):
        for j in range((n/2),elev.shape[1]-(n/2)):
            z=extract_localgrid(elev,i,j,n)
            if z.shape[0]==0: continue

            #normal equation
            A=np.matrix([[sx2,0,0],[0,sx2,0],[0,0,n]])
            B=np.matrix([[np.sum(z*x)],[np.sum(z*y)],[np.sum(z)]])
            coeff=np.linalg.solve(A,B)
            grad[i,j],asp[i,j]=compute_gradient_aspect(coeff[0],coeff[1])
            

##    fig,ax=plt.subplots(nrows=1,ncols=2,sharex=True,sharey=True)
##    plt.sca(ax[0])
##    plt.imshow(grad)
##    plt.axis('equal')
##    plt.colorbar()
##
##    plt.sca(ax[1])
##    plt.imshow(asp)
##    plt.colorbar()

    return grad,asp
    
    

def construct_arrays(path_to_file):
    #importing files from csv to array
    lith=np.genfromtxt(path_to_file+'Lith',dtype=np.int, missing_values='nan',skip_header=6)
    elev=np.genfromtxt(path_to_file+'DEM',dtype=np.float, missing_values='nan',skip_header=6)

    print lith.shape,elev.shape

    grad,asp=generate_gradient_aspect(elev,3,10)

    #exporting to csv
    np.savetxt(path_to_file+"grad_n"+str(3), np.round(grad,1), delimiter=" ", fmt='%.1f')
    np.savetxt(path_to_file+"asp_n"+str(3), np.round(asp,1), delimiter=" ", fmt='%.1f')
##    #creating x and y coordinate mesh
##    rows=lith.shape[0]
##    cols=lith.shape[1]    
##    x,y=np.meshgrid(np.arange(cols),np.arange(rows))

    #reducing array to 1D
    lith=np.ravel(lith)
    elev=np.ravel(elev)
    grad=np.ravel(grad)
    asp=np.ravel(asp)
    
    #constructing dataframe from arrays
    df=pd.DataFrame(data={'elev':elev,    
                          'Li':lith,
                          'grad':grad,
                          'asp':asp,}, index=None)
    #removing nan values
    df=df[(df.Li!=-1)]
    df.dropna(axis=0, how='any', inplace=True)

    #exporting to csv
    df.to_csv(path_to_file+"DF_el_li_gr_asp")
    print df

    return




#MAIN
path_to_file='/home/eglsais/Dropbox/Thesis/csv_files_originals/'

##construct_arrays(path_to_file)
##
##
##dx=(1,1,-1,-1)
##dy=(1,-1,1,-1)
##Q=(225,315,135,45)
##
##for i in range(4):
##
##    grad,asp=compute_gradient_aspect(dx[i],dy[i])
##    print dx[i],dy[i],grad,asp,Q[i]








